#!/bin/bash

CN_celeb1_path=/ceph/home/zhangy20/datasets/CN-Celeb1
CN_celeb2_path=~/datasets/CN-Celeb2
musan_path=~/datasets/musan
rirs_path=~/datasets/RIRS_NOISES

trials_path=data/trials.lst

nnet_type=Res2Next50_quarter # ResNet34_quarter
pooling_type=ASP
loss_type=amsoftmax

embedding_dim=256
vad=false

. ./path.sh

stage=3
echo stage $stage

# format data dir structure by soft link
if [ $stage -eq 0 ];then
	if [ ! -d data/wav_files ]; then
		mkdir -p data/wav_files
	fi

	# link and format dir structure
	rm -rf data/wav_files/dev
	mkdir -p data/wav_files/dev
	for spk in `cat ${CN_celeb1_path}/dev/dev.lst`;do
		ln -s ${CN_celeb1_path}/data/${spk} data/wav_files/dev/cnc1-$spk
	done

	for spk in `ls ${CN_celeb2_path}/data`;do
		ln -s ${CN_celeb2_path}/data/${spk} data/wav_files/dev/cnc2-$spk
	done

	echo format CN-Celeb1 trials
	python3 local/format_trials.py \
		--dataset_dir $CN_celeb1_path \
		--output_trial_path $trials_path
fi


# build data list
if [ $stage -eq 1 ]; then
	extension=wav
	if [[ $vad == true ]];then
		echo apply vad
		extension=vad
		python3 $SPEAKER_TRAINER_ROOT/scripts/vad.py --data_dir $CN_celeb1_path --num_jobs 30
		python3 $SPEAKER_TRAINER_ROOT/scripts/vad.py --data_dir $CN_celeb2_path --num_jobs 30
	fi

	echo build dev data list
	python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
		--extension $extension \
		--dataset_dir data/wav_files/dev \
		--data_list_path data/dev_list.csv

	echo build musan data list
	python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
		--extension wav \
		--dataset_dir $musan_path \
		--data_list_path data/musan_list.csv

	echo build rirs data list
	python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
		--extension wav \
		--dataset_dir $rirs_path \
		--data_list_path data/rirs_list.csv
fi


if [ $stage -eq 2 ];then
	CUDA_VISIBLE_DEVICES=2 python3 -W ignore $SPEAKER_TRAINER_ROOT/main.py \
		--nnet_type $nnet_type \
		--loss_type $loss_type \
		--batch_size 200 \
		--num_workers 80 \
		--embedding_dim $embedding_dim \
		--save_top_k 50 \
		--train_list_path data/dev_list.csv \
		--musan_list_path data/musan_list.csv \
		--rirs_list_path data/rirs_list.csv \
		--max_epochs 50 \
		--max_frames 201 --min_frames 200 \
		--learning_rate 0.01 \
		--distributed_backend dp \
		--max_seg_per_spk 500 \
		--trials_path $trials_path \
		--eval_interval -1 \
		--nPerSpeaker 1 \
		--reload_dataloaders_every_epoch \
		--gpus 1
fi


if [ $stage -eq 3 ];then
	ckpt_path=/ceph/home/zhangy20/speaker_trainer/examples/VoxCeleb/verification/exp/Res2Next50_quarter_ASP_256_amsoftmax_2021-03-05-14-43-22/epoch=49_train_loss=2.12.ckpt
	CUDA_VISIBLE_DEVICES=7 python3 -W ignore $SPEAKER_TRAINER_ROOT/main.py \
		--batch_size 64 \
		--nnet_type $nnet_type \
		--num_workers 100 \
		--train_list_path data/dev_list.csv \
		--trials_path data/trials.lst \
		--gpus 1 \
		--max_frames 401 --min_frames 400 \
		--checkpoint_path $ckpt_path \
		--evaluate
	rm -rf lightning_logs
fi

